import random
import numpy as np
from scipy.cluster.vq import kmeans

def _metric(bboxs,cluster,thr):
    r = bboxs[:,None] / cluster[None]
    x = np.minimum(r,1/r).min(2)
    best = x.max(1)
    aat = (x > 1/thr).astype(np.float64).sum(1).mean()
    bpr = (best >1/thr).astype(np.float64).mean() # ! best possible recall
    return bpr, aat

def _kmeans(wh,n,thr):
    npr = np.random

    wh = wh[(wh>2/640).any(1)]
    s = wh.std(0)
    k = kmeans(wh/s,n,iter=30)[0] * s
    if len(k) != n: k = np.sort(npr.rand(n*2)).reshape(n,2)
    
    _f, _ = _metric(wh,k,thr)
    sh = k.shape
    mp = 0.9
    _s = 0.1
    for _ in range(1000):
        v = np.ones(sh)
        while (v==1).all():
            v = ((npr.random(sh) < mp) * random.random() * npr.randn(*sh) * _s + 1).clip(0.3,3.0)
        kg = k.copy() * v
        fg, _ = _metric(wh,kg,thr)
        if fg > _f:
            _f, k = fg, kg

    return k

def autoanchors(labels,anchor_num,anchor_t=4.0,nl=3,input_shape=(640,640),by_classes=True):
    """
        labels: c,xc,yc,w,h
        anchor_num: 总的锚点数量, nl * na
    """
    assert anchor_num % nl == 0

    if by_classes:
        classes = np.unique(labels[:,0])

        clusters = []
        for c in classes:
            indices = np.where(labels[:,0]==c)
            filtered_labels = labels[indices]
            filtered_bboxes = filtered_labels[:,3:5]
            print(filtered_bboxes.shape)

            #! 初次聚类
            cluster = _kmeans(filtered_bboxes, anchor_num, anchor_t)

            #! 判断聚类效果
            bpr, aat = _metric(filtered_bboxes, cluster, anchor_t)

            while bpr <= 0.98:
                cluster = _kmeans(filtered_bboxes, anchor_num, anchor_t)
                bpr, aat = _metric(filtered_bboxes,cluster,anchor_t)

            #! 对聚类结果按面积进行排序
            r = cluster[:,0] / cluster[:,1] #! 已满足anchor_t要求
            r2 = np.maximum(r, 1/r)
            sorted_r_indices = np.argsort(r2)
            cluster = cluster[sorted_r_indices]

            area = cluster[:,0].copy() * cluster[:,1].copy()
            sorted_area_indices = np.argsort(area)
            cluster = cluster[sorted_area_indices]

            clusters.append((cluster * np.array(input_shape)).reshape([nl,-1]))

        return np.concatenate(clusters,axis=-1)
    
    else:
        filtered_bboxes = labels[:,3:5]

        #! 初次聚类
        cluster = _kmeans(filtered_bboxes,anchor_num,anchor_t)

        #! 判断聚类效果
        bpr, aat = _metric(filtered_bboxes,cluster,anchor_t)
        while bpr <= 0.98:
            cluster = _kmeans(filtered_bboxes, anchor_num, anchor_t)
            bpr, aat = _metric(filtered_bboxes,cluster,anchor_t)

        #! 对聚类结果按面积进行排序
        r = cluster[:,0] / cluster[:,1] #! 已满足anchor_t要求
        r2 = np.maximum(r, 1/r)
        sorted_r_indices = np.argsort(r2)
        cluster = cluster[sorted_r_indices]  

        area = cluster[:,0].copy() * cluster[:,1].copy()
        sorted_area_indices = np.argsort(area)
        cluster = cluster[sorted_area_indices]

        return cluster * np.array(input_shape).reshape([nl,-1])
    
def checkanchors(labels,anchors,anchor_t=4.0,input_shape=(640,640)):
    bboxes = labels[:,3:5]
    if not isinstance(anchors, np.ndarray): anchors = np.array(anchors)
    cluster = anchors.reshape([-1,2]) / np.array(input_shape)

    nl = len(anchors)
    na = len(anchors[0]//2)
    anchor_num = nl * na    

    bpr, aat = _metric(bboxes, cluster, anchor_t)
    if bpr > 0.98:
        print("The Anchors is OK")
        return anchors
    
    #! 重新测算锚点
    print("The Anchors Recomputing...")
    while bpr <= 0.98:
        cluster = _kmeans(bboxes, anchor_num, anchor_t)
        bpr, aat = _metric(bboxes, cluster, anchor_t)
    print("The Anchors Recomputed")

    #! 对聚类结果按面积进行排序
    r = cluster[:,0] / cluster[:,1] #! 已满足anchor_t要求
    r2 = np.maximum(r, 1/r)
    sorted_r_indices = np.argsort(r2)
    cluster = cluster[sorted_r_indices]

    area = cluster[:,0].copy() * cluster[:,1].copy()
    sorted_area_indices = np.argsort(area)
    cluster = cluster[sorted_area_indices]

    return cluster * np.array(input_shape).reshape([nl,-1])

def _make_grid_anchors(anchors, strides):
    if not isinstance(anchors,np.ndarray): anchors = np.array(anchors)
    nl = len(strides)
    return np.stack([anchors[i]/strides[i] for i in range(nl)]).reshape([nl,-1])